import sys, glob, os, math, modelatm

#This function finds the closest value to a given one. It is used for finding the continuum points
def closest(target, collection) :
    return min((abs(target - i), i) for i in collection)[1]
    
#This function finds the difference between the Jacco and MOOG model spectra
def cont_diff(wavelength,wave,intensity,wave_abun1,intensity_abun1) :
    cont=closest(wavelength, wave)
    cont_model=closest(cont, wave_abun1)
    cont_diff=intensity[wave.index(cont)]-intensity_abun1[wave_abun1.index(cont_model)] #Jacco-MOOG is happening here. So +ve means the model is too low
    return [cont, cont_diff]
    
def avgcont(wave,intensity,wave_abun1,intensity_abun1) :
	cont_4538=cont_diff(4538.0,wave,intensity,wave_abun1,intensity_abun1)
	cont_4559=cont_diff(4559.5,wave,intensity,wave_abun1,intensity_abun1)
	cont_4567=cont_diff(4567.5,wave,intensity,wave_abun1,intensity_abun1)
	avg_cont_shift=(cont_4538[1]+cont_4559[1]+cont_4567[1])/3
	return avg_cont_shift
	
def indexfinder(start_wavelength,end_wavelength,wave,wave_abun1):
	I=[]
	#This two lines find the first index of the wavelength point in the observed vL that is greater than wavels and wavelf
	K_start = next(x[0] for x in enumerate(wave) if x[1] > int(start_wavelength))
	K_end = next(x[0] for x in enumerate(wave) if x[1] > int(end_wavelength))
	K=K_start
	while K<K_end:
		x = wave[K]
		I.append(wave_abun1.index(x))
		K=K+1
	return I,K_start,K_end
	#I is the index of the values in the synth that match the values in obs. K_start and K_end are the values within the obs that are in the range specified.

def chi2(I,K,K_end,intensity,intensity_abun1):
	#This loops through all the wavelengths between wavels and wavelf. It finds the wavelength point that corresponds in the model
	#It then subtracts the observed intensity from the model intensity and stores this value in delta_w
	delta_w=[]
	r=0
	while K<K_end:
		delta_w.append(intensity[K]-intensity_abun1[I[r]])
		K=K+1
		r=r+1
	#This loop calculutes out chisquare. It finds the length of delta_w and substracts four since you cannot measture the chi
	#for the first two and last two wavelength points
	sigma_N=[]
	for w in range(len(delta_w)-4):
		w=w+2
		sigma_N.append(((delta_w[w+2]+delta_w[w+1]-delta_w[w-1]+delta_w[w-2])/2)**2)
	return sum(sigma_N)

def moogmodelcreator_atmosphy(temp,gravity,FeH,abundance,resolution,wavels,wavelf):
	metallicity=str(math.trunc(FeH*-10))
	specfile=glob.glob("/home/jeffrey/PHD/ruawahia/data/synspec/%(res)s/%(s)s_%(f)s/t%(temp)03dg%(grav)1.2fm%(FeH)1.2fk2.00a0.00Ba%(abundance)1.1f.spec" \
	% {"res": resolution,"s": wavels, "f":wavelf, "temp": temp, "grav": gravity, "FeH": FeH, "abundance":abundance})
	if specfile==[]:	
		#Runs MOOG and gives an error if MOOG doesn't finish properly. MOOG has been modified to exit automatically
		os.chdir('/home/jeffrey/PHD/software/MOOG2010')
		cmd='./MOOG_nokey'
		moog = os.system(cmd)
		if moog:
			print 'running MOOG failed!!!11one!!'; sys.exit(1)
		modelatm.cutandmove(temp,gravity,FeH,resolution,wavels,wavelf,abundance)

def moogmodelcreator(temp,gravity,metallicity,resolution,wavels,wavelf):
	FeH=(float(metallicity)*-1)/10
	for abundance in (0.0,+0.5,+1.0):
		filelist=glob.glob('/home/jeffrey/PHD/ruawahia/data/synspec/T'+str(temp)+'/g'+str(gravity)+'/fem'+metallicity+'/res'+resolution+'/'+wavels+'_'+wavelf+'_'+str(abundance)+'.spec')
		if filelist==[]:
			try:
				os.makedirs('/home/jeffrey/PHD/ruawahia/data/synspec/T'+str(temp)+'/g'+str(gravity)+'/fem'+metallicity+'/res'+resolution)
			except (OSError):
					pass#outputfile.write("Already got this directory\n")
			#outputfile.write('File does not exist, creating it using MOOG\n')
			modelatm.modelmaker(metallicity, FeH, temp, gravity)
			modelatm.batchfilemaker(wavels, wavelf, abundance, resolution)
			#Runs MOOG and gives an error if MOOG doesn't finish properly. MOOG has been modified to exit automatically
			os.chdir('/home/jeffrey/PHD/software/MOOG2010')
			cmd='./MOOG_nokey'
			moog = os.system(cmd)
			if moog:
				print 'running MOOG failed!!!11one!!'; sys.exit(1)
			modelatm.cutandmove(temp,gravity,metallicity,resolution,wavels,wavelf,abundance)
	
#Loads the spectrum file. This has to deal with two spectrum files if present
def load_vanloon_spectrum(nS,specFile):
	words=specFile.split()
	ifile = open('/home/jeffrey/PHD/ruawahia/data/jaccodata/'+words[nS-1],'r')
	wave=[]; intensity=[]
	for line in ifile:
	    quadruple = line.split()
	    if quadruple[1]=="*********":
			continue
	    wave.append(float(quadruple[0])); intensity.append(float(quadruple[2]))
	return wave, intensity, words[nS-1]
	
def load_model_spectrum(temp,gravity,metallicity,resolution,wavels,wavelf,abundance):
	#Loading the model spectra
	moogmodelcreator(temp,gravity,metallicity,resolution,wavels,wavelf)
	synthesis = open("/home/jeffrey/PHD/ruawahia/data/synspec/T"+str(temp)+"/g"+str(gravity)+"/fem"+metallicity+"/res"+str(resolution)+"/"+wavels+"_"+wavelf+"_"+str(abundance)+".spec")
	wave_abun1=[]; intensity_abun1=[];
	for line in synthesis:
	    double=line.split()
	    wave_abun1.append(float(double[0]))
	    intensity_abun1.append(float(double[1]))
	return wave_abun1,intensity_abun1
	
def load_model_spectrum_atmosphy(temp,gravity,FeH,resolution,wavels,wavelf,abundance):
	#Loading the model spectra
	synthesis = open("/home/jeffrey/PHD/ruawahia/data/synspec/res%(res)s/%(s)s_%(f)s/t%(temp)03dg%(grav)1.2fm%(FeH)1.2fk2.00a0.00Ba%(abundance)1.1f.spec" \
		% {"res": resolution,"s": wavels, "f":wavelf, "temp": temp, "grav": gravity, "FeH": FeH, "abundance":abundance})
	wave_abun1=[]; intensity_abun1=[];
	for line in synthesis:
	    double=line.split()
	    wave_abun1.append(float(double[0]))
	    intensity_abun1.append(float(double[1]))
	return wave_abun1,intensity_abun1
	
def eq3(BmV,FeH_me):
	sigma_eff=0.5716+BmV*0.5404+BmV**2*(-6.126*10**(-2))+-4.862*10**(-2)*BmV*FeH_me+-1.777*10**(-2)*FeH_me+-7.969*10**(-3)*FeH_me**2
	Teff=5040/sigma_eff
	return Teff
	
def eq4(BmV,FeH_me):
	sigma_eff=0.6177+BmV*0.4354+BmV**2*-4.025*10**(-3)+5.204*10**(-2)*BmV*FeH_me+-0.1127*FeH_me+-1.385*10**(-2)*FeH_me**2
	Teff=5040/sigma_eff
	return Teff

def Alonso_BmV(BmV,FeH_me):
	#Now to calculate the temperatures using Alsonso with B-V
	temp_BmV="";
	BmV_spread_eq3=[0.2,0.35,0.35,0.5]; BmV_spread_eq4=[1.9,1.8,1.35,1.0]
	FeH_low=[0.2,-0.5,-1.5,-2.5]; FeH_high=[-0.5,-1.5,-2.5,-3.0]
	#In the cross over region, Alonso recommends using linear interpolation for (B-V)
	if BmV>=0.7 and BmV<=0.8:
		m=(eq4(0.8,FeH_me)-eq3(0.7,FeH_me))/0.1
		c=eq3(0.7,FeH_me)-m*0.7
		temp_BmV=m*BmV+c
	else:		
		for l in range(4):
			if BmV>=BmV_spread_eq3[l] and BmV<=0.8 and FeH_me<=FeH_low[l] and FeH_me>FeH_high[l]:
				temp_BmV=eq3(BmV,FeH_me)
				continue
		for l in range(4):
			if BmV>=0.7 and BmV<=BmV_spread_eq4[l] and FeH_me<=FeH_low[l] and FeH_me>FeH_high[l]:
				temp_BmV=eq4(BmV,FeH_me)
				continue
	return temp_BmV
	
def Alonso_BmV_dwarf(BmV,FeH_me):
	#Now to calculate the temperatures using Alsonso with V-K
	temp_BmV="";
	BmV_min=[0.2,0.3,0.35,0.3]; BmV_max=[1.5,1.0,0.9,0.8]
	FeH_low=[0.2,-0.5,-1.5,-2.5]; FeH_high=[-0.5,-1.5,-2.5,-3.0]
	for l in range(4):
		if BmV>=BmV_min[l] and BmV<=BmV_max[l] and FeH_me<=FeH_low[l] and FeH_me>FeH_high[l]:
			sigma_eff=0.541+0.533*BmV+0.007*BmV**2-0.019*BmV*FeH_me-0.047*FeH_me-0.011*FeH_me**2
			temp_BmV=5040/sigma_eff
			continue
	if temp_BmV<1:
		temp_BmV=""
	return temp_BmV
	
def TCS_transforms_unfixed(J_2MASS,K_2MASS,V,VmK_reddening,JmK_extinction):
	JmK_ESO=((J_2MASS-K_2MASS)+0.008)/0.956
	K_ESO=K_2MASS-0.005*JmK_ESO+0.045
	K_TCS=K_ESO-0.042+0.006*JmK_ESO
	JmK_TCS=-0.012+0.910*JmK_ESO
	K=K_TCS+0.048+0.014*JmK_TCS
	VmK_TCS=0.050+0.993*(V-K)-VmK_reddening
	JmK_TCS=JmK_TCS-JmK_extinction
	return [JmK_TCS,VmK_TCS]
	
def TCS_transforms(J_2MASS,K_2MASS,V):
	JmK_ESO=((J_2MASS-K_2MASS)+0.008)/0.956
	K_ESO=K_2MASS-0.005*JmK_ESO+0.045
	K_TCS=K_ESO-0.042+0.006*JmK_ESO
	JmK_TCS=-0.012+0.910*JmK_ESO
	K=K_TCS+0.048+0.014*JmK_TCS
	VmK_TCS=0.050+0.993*(V-K)-0.3  #this adds in the interstellar reddening correction as given in Johnson et al (2008)
	JmK_TCS=JmK_TCS-0.06 #Harris (1996) reddening
	return [JmK_TCS,VmK_TCS]



def Alonso_VmK(VmK_TCS,FeH_me):
	#Now to calculate the temperatures using Alsonso with V-K
	Teff_1=0; Teff_2=0; temp_VmK="";
	VmK_spread_eq8=[0.2,1.0,1.2,1.7]; VmK_spread_eq9=[4.9,4.6,3.4,2.8]
	FeH_low=[0.2,-0.5,-1.5,-2.5]; FeH_high=[-0.5,-1.5,-2.5,-3.0]
	for l in range(4):
		if VmK_TCS>=VmK_spread_eq8[l] and VmK_TCS<=2.5 and FeH_me<=FeH_low[l] and FeH_me>FeH_high[l]:
			sigma_eff=0.5558+VmK_TCS*0.2105+VmK_TCS**2*1.981*10**(-3)+-9.965*10**(-3)*VmK_TCS*FeH_me+1.325*10**(-2)*FeH_me+-2.726*10**(-3)*FeH_me**2
			Teff_1=5040/sigma_eff
			continue
	for l in range(4):
		if VmK_TCS>=2.0 and VmK_TCS<=VmK_spread_eq9[l] and FeH_me<=FeH_low[l] and FeH_me>FeH_high[l]:
			sigma_eff=0.3770+VmK_TCS*0.3660+VmK_TCS**2*-3.170*10**(-2)+-3.074*10**(-3)*VmK_TCS*FeH_me+-2.765*10**(-3)*FeH_me+-2.973*10**(-3)*FeH_me**2
			Teff_2=5040/sigma_eff
			continue
	if Teff_1<1 or Teff_2<1:
		temp_VmK=max(Teff_1,Teff_2)
	else:
		temp_VmK=(Teff_1+Teff_2)/2
	if temp_VmK<1:
		temp_VmK=""
	return temp_VmK
	
def Alonso_VmK_dwarf(VmK_TCS,FeH_me):
	#Now to calculate the temperatures using Alsonso with V-K
	Teff_1=0; Teff_2=0; temp_VmK="";
	VmK_spread_eq7=[0.4,0.8,1.1,1.1]; VmK_spread_eq8=[4.1,3.0,2.4,2.2]
	FeH_low=[0.2,-0.5,-1.5,-2.5]; FeH_high=[-0.5,-1.5,-2.5,-3.5]
	for l in range(4):
		if VmK_TCS>=VmK_spread_eq7[l] and VmK_TCS<=1.6 and FeH_me<=FeH_low[l] and FeH_me>FeH_high[l]:
			sigma_eff=0.555+0.195*VmK_TCS+0.013*VmK_TCS**2-0.008*VmK_TCS*FeH_me+0.009*FeH_me-0.002*FeH_me**2
			Teff_1=5040/sigma_eff
			continue
	for l in range(4):
		if VmK_TCS>=1.6 and VmK_TCS<=VmK_spread_eq8[l] and FeH_me<=FeH_low[l] and FeH_me>FeH_high[l]:
			sigma_eff=0.566+0.217*VmK_TCS-0.003*VmK_TCS**2-0.024*VmK_TCS*FeH_me+0.037*FeH_me-0.002*FeH_me**2
			Teff_2=5040/sigma_eff
			continue
	if Teff_1<1 or Teff_2<1:
		temp_VmK=max(Teff_1,Teff_2)
	else:
		temp_VmK=(Teff_1+Teff_2)/2
	if temp_VmK<1:
		temp_VmK=""
	return temp_VmK
	
def Alonso_JmK(JmK_TCS,FeH_me):
	#Now to calculate the temperatures using Alsonso with V-K
	Teff_1=0; Teff_2=0; temp_JmK="";
	JmK_min=[0.0,0.2,0.3,0.4]; JmK_max=[1.1,1.0,0.9,0.8]
	FeH_low=[0.2,-0.5,-1.5,-2.5]; FeH_high=[-0.5,-1.5,-2.5,-3.0]
	for l in range(4):
		if JmK_TCS>=JmK_min[l] and JmK_TCS<=JmK_max[l] and FeH_me<=FeH_low[l] and FeH_me>FeH_high[l]:
			sigma_eff=0.5816+0.9134*JmK_TCS+-0.1443*JmK_TCS**2
			temp_JmK=5040/sigma_eff
			continue
	if temp_JmK<1:
		temp_JmK=""
	return temp_JmK
	
def Alonso_JmK_dwarf(JmK_TCS,FeH_me):
	#Now to calculate the temperatures using Alsonso with V-K
	Teff_1=0; Teff_2=0; temp_JmK="";
	JmK_min=[0.05,0.15,0.25,0.20]; JmK_max=[0.85,0.65,0.75,0.60]
	FeH_low=[0.2,-0.5,-1.5,-2.5]; FeH_high=[-0.5,-1.5,-2.5,-3.5]
	for l in range(4):
		if JmK_TCS>=JmK_min[l] and JmK_TCS<=JmK_max[l] and FeH_me<=FeH_low[l] and FeH_me>FeH_high[l]:
			sigma_eff=0.582+0.799*JmK_TCS+0.085*JmK_TCS**2
			temp_JmK=5040/sigma_eff
			continue
	if temp_JmK<1:
		temp_JmK=""
	return temp_JmK

def Alonso_logg(temp,FeH_me,M_V):
	log_g_sun=4.44 
	T_sun=5777
	M_star=0.8
	BmV_sun=0.656
	M_bol_sun=4.75 
	#Finally calculate the surface gravity and the microturbulence
	log_temp=math.log10(temp)
	X=log_temp-3.52
	Teff_spread_eq17=[3.5,3.56,3.58,3.61]
	Teff_spread_eq18=[3.96,3.83,3.80,3.74]
	FeH_low=[0.2,-0.5,-1.5,-2.5]; FeH_high=[-0.5,-1.5,-2.5,-3.0]
	BC_1=99; BC_2=99; BC=99
	for l in range(4):
		if log_temp>=Teff_spread_eq17[l] and log_temp<=3.67 and FeH_me<=FeH_low[l] and FeH_me>FeH_high[l]:
			BC_1=(-5.531*10**-2)/X-0.6177+4.420*X-2.669*X**2+0.6943*X*FeH_me-0.1071*FeH_me-8.612*10**-3*FeH_me**2
			continue
	for l in range(4): 	
		if log_temp>=3.65 and log_temp<=Teff_spread_eq18[l] and FeH_me<=FeH_low[l] and FeH_me>FeH_high[l]:
			BC_2=(-9.930*10**-2)/X+2.887*10**-2+2.275*X-4.425*X**2+0.3505*X*FeH_me-5.558*10**-2*FeH_me-5.375*10**-3*FeH_me**2
			continue
	if BC_1>50 or BC_2>50:
		BC=min(BC_1,BC_2)
		M_bol=M_V+BC
		log_g=0.40*(M_bol-M_bol_sun)+log_g_sun+4*math.log10(temp/T_sun)+math.log10(M_star)
	else:
		BC=(BC_1+BC_2)/2
		M_bol=M_V+BC
		log_g=0.40*(M_bol-M_bol_sun)+log_g_sun+4*math.log10(temp/T_sun)+math.log10(M_star)
	if BC>50:
		log_g=""
		
	return log_g
	
def spectrum_fitting_prelim(nS,spFile,resolution,wavels,wavelf):
	[wave,intensity,specfile]=load_vanloon_spectrum(nS,spFile)
	[wave_abun1,intensity_abun1]=load_model_spectrum_atmosphy(5090,2.6,-1.9,resolution,wavels,wavelf,0.5)
	[I,K_start,K_end]=indexfinder(4510,4590,wave,wave_abun1)
	return(wave,intensity,I,K_start,K_end,specfile)
	
def spectrum_fitter_fixed_temp(temp,resolution,wavels,wavelf,nS,spFile,FeH_me):
	best_chi=99; best_parms=[0,0,0,0]; temp_round=int(round(temp/250)*250)
	[wave,intensity,I,K_start,K_end,specfile]=spectrum_fitting_prelim(nS,spFile,resolution,wavels,wavelf)
	if FeH_me==-1.49999:
		FeH_range=[-0.5, -1.0, -1.5, -2.0, -2.5]
	else:
		FeH_me=round(FeH_me/0.5)*0.5
		FeH_range=[FeH_me-0.5,FeH_me,FeH_me+0.5]
	for grav_round  in [0.5, 1.0, 1.5, 2.0, 2.5, 3.0]:
		for FeH in [-0.5, -1.0, -1.5, -2.0, -2.5]:
			metallicity=str(math.trunc(FeH*-10))
			[wave_abun1,intensity_abun1]=load_model_spectrum(temp_round,grav_round,metallicity,resolution,wavels,wavelf,'0.0')
			avg_cont_shift=avgcont(wave,intensity,wave_abun1,intensity_abun1)
			for abundance in [0.0, 0.5, 1.0]:
				[wave_abun1,intensity_abun1]=load_model_spectrum(temp_round,grav_round,metallicity,resolution,wavels,wavelf,abundance)
				for index, object in enumerate(intensity_abun1):
				    intensity_abun1[index] = object+avg_cont_shift
				chi2_1=chi2(I,K_start,K_end,intensity,intensity_abun1)
				if chi2_1<best_chi:
					best_chi=chi2_1
					best_parms = [temp_round,grav_round,FeH,abundance,chi2_1,specfile]
	return best_parms
	
def spectrum_fitter_atmosphy(temp,gravity,resolution,wavels,wavelf,nS,spFile):
	best_chi=99; best_parms=[0,0,0,0];
	[wave,intensity,I,K_start,K_end,specfile]=spectrum_fitting_prelim(nS,spFile,resolution,wavels,wavelf)
	for FeH in range(20):
		FeH=float((FeH+3))/-10; metallicity=str(math.trunc(FeH*-10))
		[wave_abun1,intensity_abun1]=load_model_spectrum_atmosphy(temp,gravity,FeH,resolution,wavels,wavelf,0.5)
		avg_cont_shift=avgcont(wave,intensity,wave_abun1,intensity_abun1)
		for abundance in [0.0, 0.5, 1.0]:
			[wave_abun1,intensity_abun1]=load_model_spectrum_atmosphy(temp,gravity,FeH,resolution,wavels,wavelf,abundance)
			for index, object in enumerate(intensity_abun1):
			    intensity_abun1[index] = object+avg_cont_shift
			chi2_1=chi2(I,K_start,K_end,intensity,intensity_abun1)
			if chi2_1<best_chi:
				best_chi=chi2_1
				best_parms = [temp,gravity,FeH,abundance,chi2_1,specfile]
	return best_parms
	
def spectrum_fitter_floating_temp(temp,resolution,wavels,wavelf,nS,spFile):
	best_chi=99; best_parms=[0,0,0,0]; temp_round=int(round(temp/250)*250)
	[wave,intensity,I,K_start,K_end,specfile]=spectrum_fitting_prelim(nS,spFile,resolution,wavels,wavelf)
	temp_middle=int(round(temp/250)*250)
	for temp_round in [temp_middle-250,temp_middle,temp_middle+250]:
		for grav_round in [0.5, 1.0, 1.5, 2.0, 2.5, 3.0]:
			for FeH in [-0.5, -1.0, -1.5, -2.0, -2.5]:
				metallicity=str(math.trunc(FeH*-10))
				[wave_abun1,intensity_abun1]=load_model_spectrum(temp_round,grav_round,metallicity,resolution,wavels,wavelf,'0.0')
				avg_cont_shift=avgcont(wave,intensity,wave_abun1,intensity_abun1)
				for abundance in [0.0, 0.5, 1.0]:
					[wave_abun1,intensity_abun1]=load_model_spectrum(temp_round,grav_round,metallicity,resolution,wavels,wavelf,abundance)
					for index, object in enumerate(intensity_abun1):
					    intensity_abun1[index] = object+avg_cont_shift
					chi2_1=chi2(I,K_start,K_end,intensity,intensity_abun1)
					if chi2_1<best_chi:
						best_chi=chi2_1
						best_parms = [temp_round,grav_round,FeH,abundance,chi2_1,specfile]
	return best_parms

def FeH_from_colours(VmK_TCS,BmV):
	FeH_me_best=0.0
	temp_diff_best=1000.0
	for FeH_me in range(210):
		FeH_me=float(FeH_me)/-100
		temp_VmK=Alonso_VmK(VmK_TCS,FeH_me)
		temp_BmV=Alonso_BmV(BmV,FeH_me)
		if temp_BmV=="" or temp_VmK=="":
			continue
		temp_diff=abs(temp_VmK-temp_BmV)
		#if temp_diff<1:
			#FeH_me_best=FeH_me
			#break
		if temp_diff_best>temp_diff:
			temp_diff_best=temp_diff
			FeH_me_best=FeH_me
	temp_VmK=Alonso_VmK(VmK_TCS,FeH_me_best)
	temp_BmV=Alonso_BmV(BmV,FeH_me_best)
	return temp_VmK, temp_BmV, FeH_me_best
